import ExifReader from 'exifreader';
import text from 'png-chunk-text';
import extractChunks from 'png-chunks-extract';
import prettyBytes from 'pretty-bytes';
import { useEffect, useState } from 'react';

const readNovelAITag = async (fileOrBlob: Blob | File) => {
  const buf = await fileOrBlob.arrayBuffer();
  let chunks = [];
  try {
    chunks = extractChunks(new Uint8Array(buf));
  } catch (err) {
    return chunks;
  }
  const textChunks = chunks
    .filter(function (chunk) {
      return chunk.name === 'tEXt' || chunk.name === 'iTXt';
    })
    .map(function (chunk) {
      if (chunk.name === 'iTXt') {
        const data = chunk.data.filter((x) => x != 0x0);
        const txt = new TextDecoder().decode(data);
        return {
          keyword: '信息',
          text: txt.slice(10),
        };
      }
      return text.decode(chunk.data);
    });
  return textChunks;
};
const handleWebUiTag = (data) => {
  const [prompts, otherParas] = data.text.split('Steps: ');
  const promptSplit = prompts.split('Negative prompt: ');
  const negativePrompt = promptSplit.length > 1 ? promptSplit[1] : '无';
  return [
    {
      keyword: '提示词',
      text: promptSplit[0],
    },
    {
      keyword: '负面提示词',
      text: negativePrompt,
    },
    {
      keyword: '其他参数',
      text: `Steps: ${otherParas}`,
    },
  ];
};
export async function readFileInfo(file: File | Blob) {
  let nai = await readNovelAITag(file);
  if (nai.length == 1) {
    nai = handleWebUiTag(nai[0]);
  }
  const ok = [
    { key: '文件名', value: file.name },
    { key: '文件大小', value: prettyBytes(file.size) },
    ...nai.map((v, k) => {
      return {
        key: v.keyword,
        value: v.text,
      };
    }),
  ];
  if (nai.length == 0) {
    ok.push({
      key: '提示',
      value: '这可能不是一张 Stable Diffusion 生成的图或者不是原图, 经过了压缩',
    });
  }
  return ok;
}
const readExif = async (file) => {
  const data = await ExifReader.load(file);
  const entries = Object.entries(data);
  return entries.map(([key, value]) => ({ key, value }));
};

const useInspactImage = (fileOrBlob: File | Blob | null) => {
  const [imageInfo, setImageInfo] = useState<{ key: string; value: string }[]>(
    []
  );
  const [exif, setExifRef] = useState<{ key: string; value: any }[]>([]);

  useEffect(() => {
    (async () => {
      if (fileOrBlob) {
        // const _exif = await readExif(fileOrBlob);
        // setExifRef(_exif);
        const imgfileInfo = await readFileInfo(fileOrBlob);
        setImageInfo(imgfileInfo);
      } else {
        setExifRef([]);
        setImageInfo([]);
      }
    })();
  }, [fileOrBlob]);

  return { imageInfo, exif };
};

export default useInspactImage;
